/**
 * @file linspaced_custom.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "kernel_operator.h"
constexpr uint32_t BLOCK_DIM = 8;                        // 核心数
constexpr uint32_t ALIGN_ELEM = 8;                       // 32B对齐
constexpr uint32_t BUFFER_NUM = 2;                       // 双缓冲
constexpr uint32_t ALIGN_BASE = BLOCK_DIM * ALIGN_ELEM;  // 对齐基准 = 核心数 * 单tile元素数
/*
constexpr uint32_t SIG_HALF = 0;
constexpr uint32_t SIG_FLOAT = 1;
constexpr uint32_t SIG_INT16 = 2;
constexpr uint32_t SIG_INT32 = 3;
*/
template <typename T>
class KernelLinSpaceD {
public:
    __aicore__ inline KernelLinSpaceD() {}

    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, const uint32_t n, GM_ADDR z)
    {
        const T startT = *reinterpret_cast<const __gm__ T*>(x);
        const T endT = *reinterpret_cast<const __gm__ T*>(y);
        this->start = static_cast<float>(startT);
        this->end = static_cast<float>(endT);
        this->origSize = n;  
        this->step = (origSize > 1) ? (end - start) / (origSize - 1) : 0.0f;  
        this->totalAlignedSize = ((n + ALIGN_BASE - 1) / ALIGN_BASE) * ALIGN_BASE;
        uint32_t coreId = AscendC::GetBlockIdx();  
        this->blockSize = totalAlignedSize / BLOCK_DIM;  
        this->blockOffset = coreId * this->blockSize;       
        this->actualTileNum = this->blockSize / ALIGN_ELEM; 
        int32_t maxTileSize = ALIGN_ELEM;  

        zGm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(z) + this->blockOffset, this->blockSize);
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, maxTileSize * sizeof(float));
    }

    __aicore__ inline void Process()
    {
        if (origSize == 1) 
        {
            AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
            for (int32_t i = 0; i < ALIGN_ELEM; i++) tileOutput.SetValue(i, (i == 0) ? start : 0.0f);
            outQueueZ.EnQue(tileOutput);
            CopyOut(0); 
            return;
        }
        for (int32_t tileIdx = 0; tileIdx < this->actualTileNum; tileIdx++)
        {
            Compute(tileIdx);
            CopyOut(tileIdx);
        }
    }

private:
    __aicore__ inline void CopyIn() {}
    
    __aicore__ inline void Compute(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> tileOutput = outQueueZ.AllocTensor<float>();
        int32_t tileSize = ALIGN_ELEM; 
        int32_t tileGlobalStart = blockOffset + tileIdx * tileSize;

        #pragma unroll  
        for (int32_t i = 0; i < tileSize; i++)
        {
            int32_t globalIdx = tileGlobalStart + i;
            if (globalIdx < origSize) 
                tileOutput.SetValue(i, start + step * static_cast<float>(globalIdx));
            else
                tileOutput.SetValue(i, 0.0f);
        }
        outQueueZ.EnQue(tileOutput);
    }

    __aicore__ inline void CopyOut(int32_t tileIdx)
    {
        AscendC::LocalTensor<float> zLocal = outQueueZ.DeQue<float>();
        int32_t tileSize = ALIGN_ELEM;  
        int32_t copyOffset = tileIdx * tileSize; 
        AscendC::DataCopy(zGm[copyOffset], zLocal, tileSize);
        outQueueZ.FreeTensor(zLocal);
    }

private:
    AscendC::TPipe pipe;
    AscendC::GlobalTensor<float> zGm;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;   
    uint32_t blockSize;               // 对齐后的核心处理元素数（8的倍数）
    uint32_t blockOffset;             // 核心的全局偏移
    uint32_t origSize = 0;            // 原始有效元素数
    uint32_t totalAlignedSize = 0;    // 全局对齐元素数（ALIGN_BASE的倍数）
    uint32_t actualTileNum = 0;       // 核心的分块数（整数）
    float step = 0;
    float start = 0;
    float end = 0;   
};

extern "C" __global__ __aicore__ void linspaced_custom(GM_ADDR start, 
                                                       GM_ADDR end, 
                                                       const uint32_t n, 
                                                       GM_ADDR z){
    KernelLinSpaceD<float> op;
    op.Init(start, end, n, z);
    op.Process();
}

#ifndef ASCENDC_CPU_DEBUG 
void linspaced_custom_do(uint32_t blockDim, 
                         void *stream, 
                         uint8_t *start, 
                         uint8_t *end, 
                         const uint32_t n, 
                         uint8_t *z) { 
    linspaced_custom<<<blockDim, nullptr, stream>>>(start, end, n, z); 
} 
#endif
